package assignment04;


import gnu.trove.iterator.TIntIterator;

import java.util.List;

import nlp.util.Indexer;
import assignment04.WordAlignmentTester.Alignment;
import assignment04.WordAlignmentTester.SentencePair;
import assignment04.WordAlignmentTester.WordAligner;

public class IBMOneAligner implements WordAligner {
	private final static int NULL_INDEX = -1;
	private final static double PNULL = .2;
	private final static double PREST = .8;
	
	public static class IBMOneFactory implements WordAlignerFactory {
		IntCounter frenchCounter;
		IntCounter englishCounter;
		IntCounterMap counterMap;
		IntCounterMap tMap;
		Indexer<String> englishIndexer;
		Indexer<String> frenchIndexer;
		
		public WordAligner trainAligner(List<SentencePair> trainingSentencePairs) {
			buildIndexerAndCounter(trainingSentencePairs);
//			normalizeCounter();
			for (int i = 0; i < 5; i++)
			{
				System.err.println("model one training loop " + i);
				train(trainingSentencePairs);
				tMap = IntCounters.conditionalNormalize(counterMap);
				counterMap.reset();
			}
			return new IBMOneAligner(frenchCounter, englishCounter, tMap, englishIndexer, frenchIndexer);
		}
		
		private void buildIndexerAndCounter(List<SentencePair> trainingSentencePairs) {
			for (SentencePair pair : trainingSentencePairs) {
				List<String> eList = pair.englishWords;
				List<String> fList = pair.frenchWords;

				for (String eWord : eList) {
					englishCounter.incrementCount(englishIndexer.addAndGetIndex(eWord), 1.0);
				}
				
				for (String fWord : fList) {
					int fIndex = frenchIndexer.addAndGetIndex(fWord);
					frenchCounter.incrementCount(fIndex, 1.0);
					counterMap.setCount(NULL_INDEX, fIndex, 0.0);
				}
				
				for (String english : eList) {
					int eIndex = englishIndexer.indexOf(english);
					for (String french : fList) {
						int fIndex = frenchIndexer.indexOf(french);
						counterMap.setCount(eIndex, fIndex, 0.0);
					}
				}
			}
			
			//normalize
			englishCounter.normalize();
			frenchCounter.normalize();
			
			System.out.println("there are " + englishCounter.size() + " english word types");
			System.out.println("there are " + frenchCounter.size() + " french word types");
			
			for (TIntIterator outerIt = counterMap.keySet().iterator(); outerIt.hasNext();) {
				int key = outerIt.next();
				IntCounter counter = counterMap.getCounter(key);
				for (TIntIterator innerIt = counter.keySet().iterator(); innerIt.hasNext();) {
					int value = innerIt.next();
					tMap.setCount(key, value, frenchCounter.getCount(value));
				}
			}
		}
		
		
		private void train(List<SentencePair> trainingSentencePairs) {
			for (SentencePair pair : trainingSentencePairs) {
				List<String> eList = pair.englishWords;
				List<String> fList = pair.frenchWords;
				innerTrain(eList, fList);
			}
		}
		
		private void innerTrain(List<String> eList, List<String> fList) {
			int numFrench = fList.size();
			int numEnglish = eList.size();
			
			for (int frenchPosition = 0; frenchPosition < numFrench; frenchPosition++) {
				int frenchIndex = frenchIndexer.indexOf(fList.get(frenchPosition));
				double tNull = tMap.getCount(NULL_INDEX, frenchIndex);
				double tSum = 0.0;
				for (int englishPosition = 0; englishPosition < numEnglish; englishPosition++) {
					int englishIndex = englishIndexer.indexOf(eList.get(englishPosition));
					tSum += tMap.getCount(englishIndex, frenchIndex);
				}
				double denorminator = (PNULL * tNull) + PREST / numEnglish * tSum;
				for (int englishPosition = 0; englishPosition < numEnglish; englishPosition++) {
					int englishIndex = englishIndexer.indexOf(eList.get(englishPosition));
					double pNormal = (PREST / numEnglish * tMap.getCount(englishIndex, frenchIndex))
							/ denorminator;
							//((PNULL * tNull) + PREST / numEnglish * tSum);
					counterMap.incrementCount(englishIndex, frenchIndex, pNormal);
				}
				double pSpecial = (PNULL * tNull)
								/ denorminator;
								//(PNULL * tNull + PREST / numEnglish * tSum);
				counterMap.incrementCount(NULL_INDEX, frenchIndex, pSpecial);
			}
		}
		
		public IBMOneFactory() {
			frenchCounter = new IntCounter();
			englishCounter = new IntCounter();
			counterMap = new IntCounterMap();
			tMap = new IntCounterMap();
			
			englishIndexer = new Indexer<String>();
			frenchIndexer = new Indexer<String>();
		}
	}
	
	
	public Alignment alignSentencePair(SentencePair sentencePair) {
		// TODO Auto-generated method stub
		Alignment alignment = new Alignment();
		int numFrenchWords = sentencePair.getFrenchWords().size();
	      int numEnglishWords = sentencePair.getEnglishWords().size();
	      for (int frenchPosition = 0; frenchPosition < numFrenchWords; frenchPosition++) {
	    	  int maxPosition = -1;
	    	  double maxScore = Double.NEGATIVE_INFINITY;
	    	  String frenchWord = sentencePair.frenchWords.get(frenchPosition);
	    	  int frenchIndex = frenchIndexer.indexOf(frenchWord);
	    	  for (int englishPosition = -1; // TODO don't forget to set to -1
	    	  englishPosition < numEnglishWords; englishPosition++) {
				int englishIndex = -1;
				if (englishPosition >= 0) {
					String englishWord = sentencePair.englishWords.get(englishPosition);
					englishIndex = englishIndexer.indexOf(englishWord);
				}
	    		  double score = tMap.getCount(englishIndex, frenchIndex);
	    		  if (score > maxScore) {
	    			  maxScore = score;
	    			  maxPosition = englishPosition;
	    		  }
	    	  }
	    	  alignment.addAlignment(maxPosition, frenchPosition, true);
	      }
	      return alignment;
	}
	
	public IBMOneAligner(IntCounter frenchCounter, IntCounter englishCounter,
			IntCounterMap tMap, Indexer<String> englishIndexer,
			Indexer<String> frenchIndexer) {
		this.frenchCounter = frenchCounter;
		this.englishCounter = englishCounter;
		this.tMap = tMap;
		this.englishIndexer = englishIndexer;
		this.frenchIndexer = frenchIndexer;

	}
	
	IntCounter frenchCounter;
	IntCounter englishCounter;
	IntCounterMap tMap;
	Indexer<String> englishIndexer;
	Indexer<String> frenchIndexer;

}
